package com.masiyi.springaiollama.controller;

import com.masiyi.springaiollama.controller.dto.AigcMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
 * @Author: masiyi
 * @Date: 2024/10/21
 * @Describe:
 */
@RestController
@RequestMapping("/ai")
public class OllamaController {
    @Autowired
    private OllamaChatModel ollamaChatModel;


    @GetMapping("/ollama")
    public Object ollamaCall(@RequestParam(value = "msg") String msg){
        return ollamaChatModel.call(new Prompt(msg));

    }
    @PostMapping(value = "/ollama-stream", produces = "text/event-stream; charset=utf-8")
    public Flux<ServerSentEvent> ollamaStream(@RequestParam(value = "msg") String msg){
        return ollamaChatModel.stream(new Prompt(msg)).map(
                message -> ServerSentEvent.builder().id("111").event("message").data(message.getResult().getOutput().getContent()).build()
        );

    }

    @PostMapping("/chat")
    public String ollamaChat(@RequestBody List<AigcMessage> aigcMessages){
        ArrayList<Message> messages = new ArrayList<>();
        for (AigcMessage aigcMessage : aigcMessages) {
            if ("user".equals(aigcMessage.role())) {
                messages.add(new UserMessage(aigcMessage.content()));
            } else if ("assistant".equals(aigcMessage.role())) {
                messages.add(new AssistantMessage(aigcMessage.content()));
            }
        }
        return ollamaChatModel.call(new Prompt(messages)).getResult().getOutput().getContent();
    }


    @PostMapping(value = "/chat-stream", produces = "text/event-stream; charset=utf-8")
    public Flux<ServerSentEvent> ollamaStream(@RequestBody List<AigcMessage> aigcMessages) {
        ArrayList<Message> messages = new ArrayList<>();
        for (AigcMessage aigcMessage : aigcMessages) {
            if ("user".equals(aigcMessage.role())) {
                messages.add(new UserMessage(aigcMessage.content()));
            } else if ("assistant".equals(aigcMessage.role())) {
                messages.add(new AssistantMessage(aigcMessage.content()));
            }
        }

        return ollamaChatModel.stream(new Prompt(messages)).map(
                message -> {
                    String content = message.getResult().getOutput().getContent()
                            .replace("\n", "\\n") // 保留换行符
                            .replace(" ", "&nbsp;"); // 保留空格
                    return ServerSentEvent.builder()
                            .id("222")
                            .event("message")
                            .data(content)
                            .build();
                }
        );
    }



}
